from typing import Dict, List
from .node_handler import NodeHandler
from ..generator import OutputGenerator, StrategyGenerator
from ..StrategiesVector import StrategiesVector
from ..shard.placement_types import DeviceMesh,AbstractTensor,OperationData, OperationDataType
from geesibling.core.types import Graph, Node

__all__ = ['OutputHandler']


class OutputHandler(NodeHandler):
    """
    A OutputHandler which deals with the sharding strategies for Output Node.
    """

    def __init__(self, node: Node, device_mesh: DeviceMesh, strategies_vector: StrategiesVector,pre_strategies_vectors:List[StrategiesVector]) -> None:
        super().__init__(node, device_mesh, strategies_vector)
        self.strategies_vector.predecessor_nodes = [node.inputs[0]]
        for index,node_name in enumerate(pre_strategies_vectors):
            if index ==0:
                self.pre_strategies_vectors = list(pre_strategies_vectors[node_name])
                break

    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        generators.append(OutputGenerator(op_data_mapping, self.device_mesh, self.pre_strategies_vectors))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        # use transposed shape for strategies
        # the strategies will be transformed back to its original shape in self.post_process
        mapping = {}
        for i in range(len(self.node.inputs)):
            input_abs_data = AbstractTensor(self.node.input_shape(i),self.node.input_type(i))
            physical_input = OperationData(name=str(self.node.inputs[i]), type=OperationDataType.INPUT, data=input_abs_data)
            mapping["input_{}".format(i)] = physical_input
        output_abs_data = AbstractTensor(self.node.output_shape(0),self.node.output_type(0))
        physical_output = OperationData(name=str(self.node), data=output_abs_data,type=OperationDataType.OUTPUT)
        mapping['output'] = physical_output
        return mapping
